#include "segment/mem_segment.h"
#include "util/file_wrapper/file_wrapper.h"

namespace hawking {
namespace indexlib {

DEFINE_uint64(dumping_threashold, 1024 * 1024 * 1024 * 8L, "building segment dump threshold");
DECLARE_int32(tomb_init_size);

bool MemSegment::ShouldDump() const {
    size_t memory_use = 0;
    for (auto iter : inverted_indexers_) {
        memory_use += iter.second->MemoryUse();
    }
    for (auto iter : attribute_indexers_) {
        memory_use += iter.second->MemoryUse();
    }
    memory_use += pk_indexer_->MemoryUse();
    memory_use += deletion_indexer_->MemoryUse();
    return memory_use > FLAGS_dumping_threashold;
}

bool MemSegment::Init(
    std::shared_ptr<IndexlibIndexConfig> schema) {
    if (unlikely(!schema)) {
        LOG(ERROR) << "Empty index config";
        return false;
    }
    schema_ = schema;
    document_creator_ = std::make_shared<DocumentCreator>(schema_);

    // pk init
    pk_indexer_ = std::make_shared<PrimaryKeyMemIndexer>();
    Status ret = pk_indexer_->Init();
    if (ret != Status::OK) {
        LOG(ERROR) << "Init PrimaryKeyMemIndexer fail with "
            << static_cast<int32_t>(ret);
        return false;
    }

    // deletion init
    deletion_indexer_ = std::make_shared<DeletionMapIndexer>();
    ret = deletion_indexer_->Init();
    if (ret != Status::OK) {
        LOG(ERROR) << "Init DeletionMapMemIndexer fail with "
            << static_cast<int32_t>(ret);
        return false;
    }

    // inverted init
    for (const auto& invert_config : schema_->inverted_configs) {
        const std::string& field_name = invert_config.field_name();
        int32_t field_id = -1;
        util::ColumnType field_type;
        if (unlikely(!GetFieldIdAndTypedByName(field_name, &field_id, &field_type))) {
            LOG(ERROR) << "Get config of field_name " << field_name
                << " fail";
            return false;
        }

        std::shared_ptr<InvertedMemIndexer> inverted_mem_indexer =
            std::make_shared<InvertedMemIndexer>();
        ret = inverted_mem_indexer->Init(field_name);
        if (ret != Status::OK) {
            LOG(ERROR) << "Init InvertedMemIndexer for field_name "
                << field_name << " fail with " << static_cast<int32_t>(ret);
            return false;
        }
        inverted_indexers_.emplace(field_id, inverted_mem_indexer);
    }

    // attribute init
    for (int32_t field_id = 0; field_id < schema_->fields_size(); ++field_id) {
        const auto& field_config = schema_->fields(field_id);
        const std::string& field_name = field_config.field_name();
        util::ColumnType field_type = field_config.field_type();

        switch (field_type) {
        case util::ColumnType::COLUMN_INT8: {
            std::shared_ptr<AttributeConvertor> attr_convertor =
                std::make_shared<SingleValueAttributeConvertor<int8_t>>(field_type);
            std::shared_ptr<AttributeMemIndexer> mem_indexer = 
                std::make_shared<SingleValueAttributeMemIndexer<int8_t>>(
                    field_id, field_type, attr_convertor);
            ret = mem_indexer->Init();
            if (ret != Status::OK) {
                LOG(ERROR) << "Init SingleValueAttributeMemIndexer for field_name "
                    << field_name << " fail with " << static_cast<int32_t>(ret);
                return false;
            }
            attribute_indexers_.emplace(field_id, mem_indexer);
        }
        break;
        case util::ColumnType::COLUMN_INT16: {
            std::shared_ptr<AttributeConvertor> attr_convertor =
                std::make_shared<SingleValueAttributeConvertor<int16_t>>(field_type);
            std::shared_ptr<AttributeMemIndexer> mem_indexer = 
                std::make_shared<SingleValueAttributeMemIndexer<int16_t>>(
                    field_id, field_type, attr_convertor);
            ret = mem_indexer->Init();
            if (ret != Status::OK) {
                LOG(ERROR) << "Init SingleValueAttributeMemIndexer for field_name "
                    << field_name << " fail with " << static_cast<int32_t>(ret);
                return false;
            }
            attribute_indexers_.emplace(field_id, mem_indexer);
        }
        break;
        case util::ColumnType::COLUMN_INT32: {
            std::shared_ptr<AttributeConvertor> attr_convertor =
                std::make_shared<SingleValueAttributeConvertor<int32_t>>(field_type);
            std::shared_ptr<AttributeMemIndexer> mem_indexer = 
                std::make_shared<SingleValueAttributeMemIndexer<int32_t>>(
                    field_id, field_type, attr_convertor);
            ret = mem_indexer->Init();
            if (ret != Status::OK) {
                LOG(ERROR) << "Init SingleValueAttributeMemIndexer for field_name "
                    << field_name << " fail with " << static_cast<int32_t>(ret);
                return false;
            }
            attribute_indexers_.emplace(field_id, mem_indexer);
        }
        break;
        case util::ColumnType::COLUMN_INT64: {
            std::shared_ptr<AttributeConvertor> attr_convertor =
                std::make_shared<SingleValueAttributeConvertor<int64_t>>(field_type);
            std::shared_ptr<AttributeMemIndexer> mem_indexer = 
                std::make_shared<SingleValueAttributeMemIndexer<int64_t>>(
                    field_id, field_type, attr_convertor);
            ret = mem_indexer->Init();
            if (ret != Status::OK) {
                LOG(ERROR) << "Init SingleValueAttributeMemIndexer for field_name "
                    << field_name << " fail with " << static_cast<int32_t>(ret);
                return false;
            }
            attribute_indexers_.emplace(field_id, mem_indexer);
        }
        break;
        case util::ColumnType::COLUMN_FLOAT: {
            std::shared_ptr<AttributeConvertor> attr_convertor =
                std::make_shared<SingleValueAttributeConvertor<float>>(field_type);
            std::shared_ptr<AttributeMemIndexer> mem_indexer = 
                std::make_shared<SingleValueAttributeMemIndexer<float>>(
                    field_id, field_type, attr_convertor);
            ret = mem_indexer->Init();
            if (ret != Status::OK) {
                LOG(ERROR) << "Init SingleValueAttributeMemIndexer for field_name "
                    << field_name << " fail with " << static_cast<int32_t>(ret);
                return false;
            }
            attribute_indexers_.emplace(field_id, mem_indexer);
        }
        break;
        case util::ColumnType::COLUMN_DOUBLE: {
            std::shared_ptr<AttributeConvertor> attr_convertor =
                std::make_shared<SingleValueAttributeConvertor<double>>(field_type);
            std::shared_ptr<AttributeMemIndexer> mem_indexer = 
                std::make_shared<SingleValueAttributeMemIndexer<double>>(
                    field_id, field_type, attr_convertor);
            ret = mem_indexer->Init();
            if (ret != Status::OK) {
                LOG(ERROR) << "Init SingleValueAttributeMemIndexer for field_name "
                    << field_name << " fail with " << static_cast<int32_t>(ret);
                return false;
            }
            attribute_indexers_.emplace(field_id, mem_indexer);
        }
        break;
        case util::ColumnType::COLUMN_STRING: {
            std::shared_ptr<AttributeConvertor> attr_convertor =
                std::make_shared<StringAttributeConvertor>(field_type);
            std::shared_ptr<AttributeMemIndexer> mem_indexer = 
                std::make_shared<MultiValueAttributeMemIndexer<char>>(
                    field_id, field_type, attr_convertor);
            ret = mem_indexer->Init();
            if (ret != Status::OK) {
                LOG(ERROR) << "Init MultiValueAttributeMemIndexer for field_name "
                    << field_name << " fail with " << static_cast<int32_t>(ret);
                return false;
            }
            attribute_indexers_.emplace(field_id, mem_indexer);
        }
        break;
        case util::ColumnType::COLUMN_INT8_LIST: {
            std::shared_ptr<AttributeConvertor> attr_convertor =
                std::make_shared<MultiValueAttributeConvertor<int8_t>>(field_type);
            std::shared_ptr<AttributeMemIndexer> mem_indexer = 
                std::make_shared<MultiValueAttributeMemIndexer<int8_t>>(
                    field_id, field_type, attr_convertor);
            ret = mem_indexer->Init();
            if (ret != Status::OK) {
                LOG(ERROR) << "Init MultiValueAttributeMemIndexer for field_name "
                    << field_name << " fail with " << static_cast<int32_t>(ret);
                return false;
            }
            attribute_indexers_.emplace(field_id, mem_indexer);
        }
        break;
        case util::ColumnType::COLUMN_INT16_LIST: {
            std::shared_ptr<AttributeConvertor> attr_convertor =
                std::make_shared<MultiValueAttributeConvertor<int16_t>>(field_type);
            std::shared_ptr<AttributeMemIndexer> mem_indexer = 
                std::make_shared<MultiValueAttributeMemIndexer<int16_t>>(
                    field_id, field_type, attr_convertor);
            ret = mem_indexer->Init();
            if (ret != Status::OK) {
                LOG(ERROR) << "Init MultiValueAttributeMemIndexer for field_name "
                    << field_name << " fail with " << static_cast<int32_t>(ret);
                return false;
            }
            attribute_indexers_.emplace(field_id, mem_indexer);
        }
        break;
        case util::ColumnType::COLUMN_INT32_LIST: {
            std::shared_ptr<AttributeConvertor> attr_convertor =
                std::make_shared<MultiValueAttributeConvertor<int32_t>>(field_type);
            std::shared_ptr<AttributeMemIndexer> mem_indexer = 
                std::make_shared<MultiValueAttributeMemIndexer<int32_t>>(
                    field_id, field_type, attr_convertor);
            ret = mem_indexer->Init();
            if (ret != Status::OK) {
                LOG(ERROR) << "Init MultiValueAttributeMemIndexer for field_name "
                    << field_name << " fail with " << static_cast<int32_t>(ret);
                return false;
            }
            attribute_indexers_.emplace(field_id, mem_indexer);
        }
        break;
        case util::ColumnType::COLUMN_INT64_LIST: {
            std::shared_ptr<AttributeConvertor> attr_convertor =
                std::make_shared<MultiValueAttributeConvertor<int64_t>>(field_type);
            std::shared_ptr<AttributeMemIndexer> mem_indexer = 
                std::make_shared<MultiValueAttributeMemIndexer<int64_t>>(
                    field_id, field_type, attr_convertor);
            ret = mem_indexer->Init();
            if (ret != Status::OK) {
                LOG(ERROR) << "Init MultiValueAttributeMemIndexer for field_name "
                    << field_name << " fail with " << static_cast<int32_t>(ret);
                return false;
            }
            attribute_indexers_.emplace(field_id, mem_indexer);
        }
        break;
        case util::ColumnType::COLUMN_FLOAT_LIST: {
            std::shared_ptr<AttributeConvertor> attr_convertor =
                std::make_shared<MultiValueAttributeConvertor<float>>(field_type);
            std::shared_ptr<AttributeMemIndexer> mem_indexer = 
                std::make_shared<MultiValueAttributeMemIndexer<float>>(
                    field_id, field_type, attr_convertor);
            ret = mem_indexer->Init();
            if (ret != Status::OK) {
                LOG(ERROR) << "Init MultiValueAttributeMemIndexer for field_name "
                    << field_name << " fail with " << static_cast<int32_t>(ret);
                return false;
            }
            attribute_indexers_.emplace(field_id, mem_indexer);
        }
        break;
        case util::ColumnType::COLUMN_DOUBLE_LIST: {
            std::shared_ptr<AttributeConvertor> attr_convertor =
                std::make_shared<MultiValueAttributeConvertor<double>>(field_type);
            std::shared_ptr<AttributeMemIndexer> mem_indexer = 
                std::make_shared<MultiValueAttributeMemIndexer<double>>(
                    field_id, field_type, attr_convertor);
            ret = mem_indexer->Init();
            if (ret != Status::OK) {
                LOG(ERROR) << "Init MultiValueAttributeMemIndexer for field_name "
                    << field_name << " fail with " << static_cast<int32_t>(ret);
                return false;
            }
            attribute_indexers_.emplace(field_id, mem_indexer);
        }
        break;
        case util::ColumnType::COLUMN_STRING_LIST: {
            std::shared_ptr<AttributeConvertor> attr_convertor =
                std::make_shared<MultiStringAttributeConvertor>(field_type);
            std::shared_ptr<AttributeMemIndexer> mem_indexer = 
                std::make_shared<MultiValueAttributeMemIndexer<std::string>>(
                    field_id, field_type, attr_convertor);
            ret = mem_indexer->Init();
            if (ret != Status::OK) {
                LOG(ERROR) << "Init MultiValueAttributeMemIndexer for field_name "
                    << field_name << " fail with " << static_cast<int32_t>(ret);
                return false;
            }
            attribute_indexers_.emplace(field_id, mem_indexer);
        }
        break;
        default:
            LOG(ERROR) << "Unknown field_type " << static_cast<int32_t>(field_type);
            return false;
        }
    }
    
    tomb_stone_ = std::make_unique<TombStone>(FLAGS_tomb_init_size);
    return true;
}

bool MemSegment::Dump_(const std::string& segment_dir) {
    std::unordered_map<DocId, DocId> old2new_docid_map =
        pk_indexer_->GenerateOldNewDocMap(deletion_indexer_);

    if (TTL()) {
        ObsoletedFilter_(&old2new_docid_map);
    }
    
    // pk dump
    std::string pk_dir =
        butil::string_printf("%s/primary_key", segment_dir.data());
    bool ret = FslibWrapper::RmDir(pk_dir) && FslibWrapper::Mkdir(pk_dir);
    if (!ret) {
        LOG(ERROR) << "Create direct " << pk_dir << " fail on dumping.";
        return false;
    }

    Status dump_ret = pk_indexer_->Dump(pk_dir, old2new_docid_map);
    if (Status::OK != dump_ret) {
        LOG(ERROR) << "Dump pk indexer to " << pk_dir
            << " fail on dumping. with error "
            << static_cast<int32_t>(dump_ret);
        return false;
    }

    // inverted dump
    std::string inverted_dir =
        butil::string_printf("%s/inverted_index", segment_dir.data());
    ret = FslibWrapper::RmDir(inverted_dir) && FslibWrapper::Mkdir(inverted_dir);
    if (!ret) {
        LOG(ERROR) << "Create direct " << inverted_dir << " fail on dumping.";
        return false;
    }
    for (auto& iter : inverted_indexers_) {
        const std::string& field_name = iter.first;
        std::string cur_inverted_dir =
            butil::string_printf("%s/%s", inverted_dir.data(), field_name.data());
        ret = FslibWrapper::RmDir(cur_inverted_dir)
            && FslibWrapper::Mkdir(cur_inverted_dir);
        if (!ret) {
            LOG(ERROR) << "Create direct " << cur_inverted_dir << " fail on dumping.";
            return false;
        }

        dump_ret = iter.second->Dump(cur_inverted_dir, old2new_docid_map);
        if (Status::OK != dump_ret) {
            LOG(ERROR) << "Dump inverted indexer to " << cur_inverted_dir
                << " fail on dumping. error with " << static_cast<int32_t>(dump_ret);
            return false;
        }
    }

    // attribute dump
    std::string attribute_dir =
        butil::string_printf("%s/attribute_index", segment_dir.data());
    ret = FslibWrapper::RmDir(attribute_dir) && FslibWrapper::Mkdir(attribute_dir);
    if (!ret) {
        LOG(ERROR) << "Create direct " << attribute_dir << " fail on dumping.";
        return false;
    }
    for (auto& iter : attribute_indexers_) {
        const std::string& field_name = iter.first;
        std::string cur_attribute_dir =
            butil::string_printf("%s/%s", attribute_dir.data(), field_name.data());
        ret = FslibWrapper::RmDir(cur_attribute_dir)
            && FslibWrapper::Mkdir(cur_attribute_dir);
        if (!ret) {
            LOG(ERROR) << "Create direct " << cur_attribute_dir << " fail on dumping.";
            return false;
        }

        dump_ret = iter.second->Dump(cur_attribute_dir, old2new_docid_map);
        if (Status::OK != dump_ret) {
            LOG(ERROR) << "Dump attribute indexer to " << cur_attribute_dir
                << " fail on dumping. error with " << static_cast<int32_t>(dump_ret);
            return false;
        }
    }

    // index config's json as meta
    std::string index_config_json = "";
    ret = json2pb::ProtoMessageToJson(*schema_, &index_config_json);
    if (!ret) {
        LOG(ERROR) << "index config ProtoMessageToJson fail, schema "
            << schema_->ShortDebugString();
        return false;
    }

    std::string meta_filename = butil::string_printf("%s/meta", segment_dir.data());
    NormalFileWriter meta_file(meta_filename);
    if (meta_file.Init() != Status::OK) {
        LOG(ERROR) << "meta file " << meta_filename << " Init fail";
        return false;
    }

    auto r = meta_file.Write(index_config_json.data(), index_config_json.length());
    if (r.Code() != Status::OK) {
        LOG(ERROR) << "meta file " << meta_filename << " Write fail";
        return false;
    }

    return true;
}

bool MemSegment::Dump(const std::string& segment_dir) {
    WaitForWriteFinish();
    LOG(INFO) << "segment " << segment_dir << " begin dump";
    bool dump_ret = Dump_(segment_dir);
    ResetSegmentState();
    return dump_ret;
}

util::HawkingRecallRetCode MemSegment::AddDocument(std::shared_ptr<NormalDocument> doc) {
    if (!CouldWrite()) {
        return false;
    }

    CallBackOnExit end_write_cb([this](){
        this->ResetSegmentState();
    });

    doc->SetDocId(doc_id_.load(std::memory_order_acquire));
    doc_id_.fetch_add(1, std::memory_order_release);
    for (auto iter : attribute_indexers_) {
        if (Status::OK != iter.second->AddDocument(doc.get())) {
            deletion_indexer_->Delete(doc->DocId());
            return false;
        }
    }

    for (auto iter : inverted_indexers_) {
        if (Status::OK != iter.second->AddDocument(doc.get())) {
            deletion_indexer_->Delete(doc->DocId());
            return false;
        }
    }

    if (Status::OK != pk_indexer_->AddDocument
            (doc->GetIndexDocument()->PrimaryKey(), doc->DocId())) {
        deletion_indexer_->Delete(doc->DocId());
        return false;
    }

    SetDirty_();
    return true;
}

}
}